package com.study.feature.extract

import org.apache.spark.ml.linalg
import org.apache.spark.ml.feature.Word2Vec
import org.apache.spark.sql.{Row, SparkSession}

/**
 * 特征提取-Word2Vec
 *
 * Word2vec是一个Estimator，它采用一系列代表文档的词语来训练word2vecmodel。
 * 该模型将每个词语映射到一个固定大小的向量。
 * word2vecmodel使用文档中每个词语的平均数来将文档转换为向量，
 * 然后这个向量可以作为预测的特征，来计算文档相似度计算等等。
 *
 * @author stephen
 * @date 2019-08-28 10:36
 */
object Word2VecDemo {

  def main(args: Array[String]): Unit = {

    val spark = SparkSession.builder()
      .appName(this.getClass.getSimpleName)
      .master("local[*]")
      .getOrCreate()

    spark.sparkContext.setLogLevel("warn")

    val documentDF = spark.createDataFrame(Seq(
      "Hi I heard about Spark".split(" "),
      "I wish Java could use case classes".split(" "),
      "Logistic regression models are neat".split(" ")
    ).map(Tuple1.apply)).toDF("text")

    val word2Vec = new Word2Vec()
      .setInputCol("text")
      .setOutputCol("result")
      // 设置特征向量的维度为3
      .setVectorSize(3)
      .setMinCount(0)
    val model = word2Vec.fit(documentDF)

    // 把文档转成特征向量
    val result = model.transform(documentDF)
    result.collect().foreach { case Row(text: Seq[_], features: linalg.Vector) =>
      println(s"Text: [${text.mkString(", ")}] => \nVector: $features\n")
    }

    spark.stop()
  }
}
